// -----------------------------
// Author: 		Dhaenens Gert	
// Creation: 	10/18/2014 1:21:49 AM				
// -----------------------------

#include "TerrainMaterial.h"

#include "Application.h"
#include "FileManager.h"
#include "RenderManager.h"
#include "Scene.h"
#include "Camera.h"

#define MAX_GEOMETRY_OUTPUT_VERTICES	0x000FFFFF

struct GeometryOutputVertex
{
	Vector4	Position;
	Vector3	Normal;
};

struct MatrixBuffer
{
	Matrix World;
	Matrix WorldViewProj;
	Matrix View;
};

struct TerrainBuffer
{
	Float	MinimumHeight;
	Float	MaximumHeight;
	Float	TerrainWidth;
	Float	TerrainHeight;
};

struct MaterialBuffer
{
	Vector3	Color;
	Float	Roughness;
	Vector4	Ambient;
	Float	Specular;
	Bool	UseTexture;
	Vector2	Padding;
};

// Tesselation Information
struct TesselationInfo
{
	Vector3	CameraPosition;
	Float	MinimumDistance;
	Float	MaximumDistance;
	Float	MaximumTessFactor;
	Vector2	Padding;
};

TerrainMaterial::TerrainMaterial(Application* lpApp)
	: Material(lpApp, RenderType::Deferred, false)
	, m_Diffuse(0.5f, 0.5f, 0.5f)
	, m_Ambient(0.2f, 0.2f, 0.2f)
	, m_Roughness(1.0f)
	, m_Specular(0.0f)
	, m_lpVertexShader(nullptr)
	, m_lpHullShader(nullptr)
	, m_lpTessInfoBuffer(nullptr)
	, m_lpDomainShader(nullptr)
	, m_lpGeometryShader(nullptr)
	, m_lpMatrixBuffer(nullptr)
	, m_lpTerrainBuffer(nullptr)
	, m_lpGeometryOutput(nullptr)
	, m_lpPixelShader(nullptr)
	, m_lpMaterialBuffer(nullptr)
	, m_lpInputLayout(nullptr)
	, m_lpDiffuseTexture(nullptr)
	, m_lpHeightTexture(nullptr)
	, m_lpSamplerState(nullptr)
{
	Init(lpApp);
}
TerrainMaterial::~TerrainMaterial()
{
	NullRelease(&m_lpSamplerState);
	NullRelease(&m_lpInputLayout);
	NullRelease(&m_lpMaterialBuffer);
	NullRelease(&m_lpPixelShader);
	NullRelease(&m_lpGeometryShader);
	NullRelease(&m_lpDomainShader);
	NullRelease(&m_lpTessInfoBuffer);
	NullRelease(&m_lpHullShader);
	NullRelease(&m_lpGeometryOutput);
	NullRelease(&m_lpTerrainBuffer);
	NullRelease(&m_lpMatrixBuffer);
	NullRelease(&m_lpVertexShader);
}

void TerrainMaterial::Init(Application* lpApp)
{
	TString vertexShader = lpApp->GetFiles()->GetFileName(_T("Shader_TerrainVertex"));
	TString hullShader = lpApp->GetFiles()->GetFileName(_T("Shader_TerrainHull"));
	TString domainShader = lpApp->GetFiles()->GetFileName(_T("Shader_TerrainDomain"));
	TString geometryShader = lpApp->GetFiles()->GetFileName(_T("Shader_TerrainGeometry"));
	TString pixelShader = lpApp->GetFiles()->GetFileName(_T("Shader_TerrainPixel"));

	m_lpVertexShader = ShaderCompiler::LoadVertexFromFile(m_lpRenderManager->GetDevice(), vertexShader.c_str());

	m_lpHullShader = ShaderCompiler::LoadHullFromFile(m_lpRenderManager->GetDevice(), hullShader.c_str());
	m_lpTessInfoBuffer = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::Constant, Usage::Default, 0, sizeof(TesselationInfo), 1, nullptr);

	m_lpDomainShader = ShaderCompiler::LoadDomainFromFile(m_lpRenderManager->GetDevice(), domainShader.c_str());

	StreamOutputElement geomElements[] =
	{
		{ 0, "WORLDPOSITION", 0, 0, 4, 0 },
		{ 0, "NORMAL", 0, 0, 3, 1 },
	};
	UInt32 geomElementCount = ARRAYSIZE(geomElements);
	m_lpGeometryShader = ShaderCompiler::LoadGeometryStreamOutputFromFile(m_lpRenderManager->GetDevice(), geometryShader.c_str(), geomElements, geomElementCount);
	m_lpMatrixBuffer = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::Constant, Usage::Default, 0, sizeof(MatrixBuffer), 1, nullptr);
	m_lpTerrainBuffer = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::Constant, Usage::Default, 0, sizeof(TerrainBuffer), 1, nullptr);
	m_lpGeometryOutput = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::StreamOutput, Usage::Default, 0, sizeof(GeometryOutputVertex), MAX_GEOMETRY_OUTPUT_VERTICES, nullptr);

	m_lpPixelShader = ShaderCompiler::LoadPixelFromFile(m_lpRenderManager->GetDevice(), pixelShader.c_str());
	m_lpMaterialBuffer = IBuffer::Create(m_lpRenderManager->GetDevice(), BufferType::Constant, Usage::Default, 0, sizeof(MaterialBuffer), 1, nullptr);

	m_lpInputLayout = IInputLayout::Create(m_lpRenderManager->GetDevice(), m_lpVertexShader);

	m_lpSamplerState = ISamplerState::Create(m_lpRenderManager->GetDevice(), AddressMode::Mirror, AddressMode::Mirror, AddressMode::Mirror);
}

void TerrainMaterial::Apply(Application* lpApp)
{
	IRenderDevice* lpDevice = m_lpRenderManager->GetDevice();

	// Bind Vertex
	lpDevice->SetVertexShader(m_lpVertexShader);

	// Bind Hull
	lpDevice->SetHullShader(m_lpHullShader);
	m_lpHullShader->SetBufferAsConstant(0, m_lpMatrixBuffer);
	m_lpHullShader->SetBufferAsConstant(1, m_lpTessInfoBuffer);
	m_lpHullShader->SetBufferAsConstant(2, m_lpTerrainBuffer);
	if (m_lpHeightTexture)
	{
		m_lpHullShader->SetTextureAsResource(0, m_lpHeightTexture.Get());
		m_lpHullShader->SetSamplerState(0, m_lpSamplerState);
	}

	// Bind Domain
	lpDevice->SetDomainShader(m_lpDomainShader);

	// Bind Geometry
	lpDevice->SetGeometryShader(m_lpGeometryShader);
	m_lpGeometryShader->SetBufferAsConstant(0, m_lpMatrixBuffer);
	m_lpGeometryShader->SetBufferAsConstant(2, m_lpTerrainBuffer);
	if (m_lpHeightTexture)
	{
		m_lpGeometryShader->SetTextureAsResource(0, m_lpHeightTexture.Get());
		m_lpGeometryShader->SetSamplerState(0, m_lpSamplerState);
	}

	// Bind Pixel
	lpDevice->SetPixelShader(m_lpPixelShader);
	m_lpPixelShader->SetBufferAsConstant(0, m_lpMaterialBuffer);
	if (m_lpDiffuseTexture)
	{
		m_lpPixelShader->SetTextureAsResource(0, m_lpDiffuseTexture.Get());
		m_lpPixelShader->SetSamplerState(0, m_lpSamplerState);
	}

	// Bind InputManager & Topology
	lpDevice->SetInputLayout(m_lpInputLayout);
	lpDevice->SetTopology(Topology::Control4PointPatchList);
}
void TerrainMaterial::UpdateMatrices(Application* lpApp)
{
	// Update the Matrix Buffer
	MatrixBuffer matrixBuffer = {};
	matrixBuffer.World = m_World;
	matrixBuffer.View = m_View;
	matrixBuffer.WorldViewProj = m_World * m_View * m_Projection;
	m_lpMatrixBuffer->Update(&matrixBuffer, sizeof(MatrixBuffer));
}
void TerrainMaterial::UpdateData(Application* lpApp)
{
	// Update Terrain Buffer
	TerrainBuffer terrBuffer = {};
	terrBuffer.MinimumHeight = m_MinimumHeight;
	terrBuffer.MaximumHeight = m_MaximumHeight;
	terrBuffer.TerrainWidth = m_TerrainWidth;
	terrBuffer.TerrainHeight = m_TerrainHeight;
	m_lpTerrainBuffer->Update(&terrBuffer, sizeof(TerrainBuffer));

	// Update Material Buffer
	MaterialBuffer matBuffer = {};
	matBuffer.Color = m_Diffuse.ToVector3();
	matBuffer.Ambient = m_Ambient.ToVector4();
	matBuffer.Specular = m_Specular;
	matBuffer.Roughness = m_Roughness;
	matBuffer.UseTexture = (m_lpDiffuseTexture) ? 1 : 0;
	m_lpMaterialBuffer->Update(&matBuffer, sizeof(MaterialBuffer));

	// Update Tesselation Buffer
	TesselationInfo tessBuffer = {};
	tessBuffer.CameraPosition = lpApp->GetScene()->GetCamera()->GetPosition();
	tessBuffer.MinimumDistance = m_MinimumDistance;
	tessBuffer.MaximumDistance = m_MaximumDistance;
	tessBuffer.MaximumTessFactor = m_MaximumTessFactor;
	m_lpTessInfoBuffer->Update(&tessBuffer, sizeof(TesselationInfo));
}

void TerrainMaterial::SetDiffuseTexture(const RefPtr<ITexture2D>& lpTexture)
{
	m_lpDiffuseTexture = lpTexture;
}
void TerrainMaterial::SetHeightTexture(const RefPtr<ITexture2D>& lpTexture)
{
	m_lpHeightTexture = lpTexture;
}

Bool TerrainMaterial::SampleGeometryOutput(const Vector3& pin, Vector3& pout, Vector3& nout) const
{
	m_lpGeometryOutput->RecalculateCount();
	UInt32 sampleSize = m_lpGeometryOutput->GetCount();
	GeometryOutputVertex* lpData = reinterpret_cast<GeometryOutputVertex*>(m_lpGeometryOutput->GetData());

	for (UInt32 i = 0; i < sampleSize / 3; ++i)
	{
		GeometryOutputVertex va = lpData[i * 3 + 0];
		GeometryOutputVertex vb = lpData[i * 3 + 0];
		GeometryOutputVertex vc = lpData[i * 3 + 0];

		Vector2 a2 = Vector2(va.Position.x, va.Position.z);
		Vector2 b2 = Vector2(vb.Position.x, vb.Position.z);
		Vector2 c2 = Vector2(vc.Position.x, vc.Position.z);

		Float denom = ((b2.y - c2.y)*(a2.x - c2.x) + (c2.x - b2.x)*(a2.y - c2.y));
		Float a = ((b2.y - c2.y)*(pin.x - c2.x) + (c2.x - b2.x)*(pin.y - c2.y)) / denom;
		Float b = ((c2.y - a2.y)*(pin.x - c2.x) + (a2.x - c2.x)*(pin.y - c2.y)) / -denom;
		Float c = 1.0f - a - b;

		Bool bInsideA = a >= 0.0f && a <= 1.0f;
		Bool bInsideB = b >= 0.0f && b <= 1.0f;
		Bool bInsideC = c >= 0.0f && c <= 1.0f;

		Bool bInside = bInsideA && bInsideB && bInsideC;
		if (!bInside)
			continue;

		// Found triangle we belong to
		Float h1 = Math::Lerp(va.Position.y, vb.Position.y, a);
		Float h2 = Math::Lerp(va.Position.y, vc.Position.y, b);
		Float h3 = Math::Lerp(vb.Position.y, vc.Position.y, c);
		Float h = (h1 + h2 + h3) / 3.0f;

		pout.x = pin.x;
		pout.y = -h;
		pout.z = pin.z;
		return true;
	}

	return false;
}